#include <algorithm>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>

#include "TreeView.h"
#include "Cats.h"

using namespace std;

// - All provisonal (until I understand why I cannot set it as a static member)
int C;
bool dNAP;
// ----------------------------------------------------------------------------

bool print_epinfo = false;

// #define DEBUG  // Comment to turn off


// ############################# CONSTRUCTORS #################################

TreeView::TreeView(TreeView* n, TreeView* r, double l)
{
	next = n;
	rev = r;
	length = l;
	name = "ROOT";
	pT = 0; // due palle
}


TreeView::TreeView(const string& treefile)
{
	pT = 0; // due palle
	ifstream tree_in(treefile.c_str());
//        string tree = "((A:1,B:2):3,C:4):0;";
//        stringbuf sb(tree, ios::in);
//        istream tree_in(&sb);
	if (!tree_in) throw string("cannot open " + treefile);
	init_structure(tree_in);
}


TreeView::TreeView(istream& tree_in)
{
	pT = 0; // due palle
	init_structure(tree_in);
}



// ######################### INITIALIZER FUNCTIONS ############################


// This function sets the tree structure while reading the file in input.
// The object on which it's called should be 'fresh of creation'... a proper
// setStructure function should also dispose of the objects that are pointed to.
//
// Normally such code would be included in a constructor, but since C++ does
// not allow a constructor to call another constructor, this would have to be
// repeated in more than one construcor.  This is why this function is defined.
//
// If tree representations are sensitive only to '(', ')' nd identifiers,
// a simpler version is possible.
//
void TreeView::init_structure(istream & in)
{
	next = this;
	int bracket_open = 0;

	char c;
	TreeView* child;

	while(in.get(c)) // exited upon return, not eof
	{
		switch (c)
		{
		case ' ' : 
		case '\n':
		case '\t':
		case '\r':
		case '\f':
		case '\v':
			break; // formatting of the tree specifier is allowed

		case '(':
			if (bracket_open++) throw string("tree construction aborted: "
				"a '(' is opened where not expected.");
			child = new TreeView(in);
			rev   = new TreeView(child, this);
			child->next = rev;
			break;

		case ',':
			if (!bracket_open) throw string("tree construction aborted: "
				"a ',' is found where not expected.");
			child = new TreeView(in);
			child->next = rev->next;
			rev->next = child;
			break;

		case ')':
			if (!bracket_open) throw string("tree construction aborted: "
				"a ')' is closed where not expected.");
			set_length(in);
			return;

		default: // everything else is assumed to be a species id:
			in.putback(c);
			get_id(in, name);
 
			if (bracket_open) throw string("tree construction aborted: "
				"taxon '" + name + "' is found where not expected.");

			rev = new TreeView(0, this);
			rev->next = rev;

			set_length(in);
			return;				
		}
	}
	throw string("tree construction aborted: end of file encountered too early"
		", perhaps a ')' was not closed.");
}

void TreeView::set_cost_benefits(const string& costfile)
{
	ifstream in(costfile.c_str());
	if (!in) throw string("cannot open " + costfile);
	set_cost_benefits(in);
}

void TreeView::set_cost_benefits(istream& in)
{
	map<string, vector<double>*> cb;
	dNAP = read_cost_benefits(in, cb, C);
	set_cost_benefits(cb);
}

void TreeView::set_cost_benefits(map<string,vector<double>*>& cb)
{
	vector<TreeView*> tips(0); // do I need that 0 to ensure it's empty?
	int n = get_tips(tips);
	while(--n >= 0) tips[n]->set_cost_benefit(cb);
}

void TreeView::set_cost_benefit(map<string,vector<double>*>& cb)
{
	pT = cb[name]; // if cb[name] was undefined, this will set it to its default, 
			       // which I assume is 0... but I may be wrong. ERR?

	if(pT == 0) throw string("tree construction aborted: '" 
		+ name + "' does not appear in the .cb file.");

	cost = 0;
	for(int c = 0; c <= C; c++)
		if((*pT)[c] > (*pT)[cost]) cost = c;
}


// ########################## MEMBER FUNCTIONS ################################

//
// Returns the probability of extinction of the clade, and transforms branch
// lengths in it.
//
double TreeView::transform_lengths_and_probs()
{
	double extinction_prob = 1;

	if (is_tip())
	{
		extinction_prob =  1 - (*pT)[0];

		if (0 < (*pT)[0] && (*pT)[0] < 1)
			for(int c = C; c >= 0; c--)  // very important: in decreasing order
				(*pT)[c] = ((*pT)[c] - (*pT)[0])/(1 - (*pT)[0]);

		// if (*pT)[0] == 1, then the length of the path to this leaf is set
		// to 0, and it doesn't really matter what the (*pT)[c] values are.
	}
	else
	{
		vector<TreeView*> c = childs();
		for(vector<TreeView*>::const_iterator i = c.begin(); i != c.end(); i++)
			extinction_prob *= (*i)->transform_lengths_and_probs();
	}
	length *= extinction_prob;

	return extinction_prob;
}

void TreeView::set_length(istream& in)
{
	char c;
	
	length = 0;

	while (in.get(c) && isspace(c));

	if (!in.fail())
	{
		if (c == ':') in >> length; // if I don't get a double, 0 should stay
		else          in.putback(c);
	}

	rev->length = length;
}

void TreeView::calc_representation()
{
	// Slightly inefficient.

	if (is_tip())
	{
		repr = "";
		repr.push_back(name[0]);
		if (name.size() > 1) repr.push_back(name[1]);
		if (name == "ROOT") repr = "ROOT";
		return;
	}
	repr = "(";
	vector<TreeView*> c = childs();
	for(vector<TreeView*>::const_iterator i = c.begin(); i != c.end(); i++)
		repr = repr + (i==c.begin() ? "" : ",") + (*i)->repr; 
	repr = repr + ")";
}

void TreeView::print_tree(ostream& out, bool semicolon)
{
	if (is_tip())
	{
		out << name << ":" << length; 
		return;
	}
	out << "(";
	
	vector<TreeView*> c = childs();
	for(vector<TreeView*>::const_iterator i = c.begin(); i != c.end(); i++)
	{
		if (i != c.begin()) out << ", ";
		(*i)->print_tree(out);
	} 
	out << "):" << length;
	
	if (semicolon) out << ";" << endl;
}


// If the TreeView in input is multifurcating at the root, that multifurcation
// is resolved into a number of bifurcations.  Note that none of the
// multifucations further below in the tree are affected.
//
// At the moment the implementation runs in O(d^2)... which is quite bad.
// With a prev pointer, it can be made O(d), because there's no need to run
// childs() to get c.back().
//
void TreeView::resolve_multifurcation()
{
	vector<TreeView*> c = childs();
	if (c.size() <= 2) return;

	TreeView* newchild = c[0]->next = new TreeView(rev, 0);
	newchild->rev = new TreeView(c[1], newchild);
	c.back()->next = newchild->rev;
	
	newchild->resolve_multifurcation();
}

// If the TreeView only has one child, this function removes that monofurcation
// at the root. Note that the resulting TreeView may still be monofurcating. To
// entirely remove monofurcations from a TreeView run 
// tv.bottom_up(&TreeView::resolve_monofurcation); (top_down would not work).
//
// If the only child is a tip, this may lose some information 
// (such as pT, or the repr)... so launch it before you set that info.
//
void TreeView::resolve_monofurcation()
{
	if (is_tip() || rev->next->next != rev) return;

	TreeView* only_child = rev->next;

	length += only_child->length;
	rev = only_child->rev;
	rev->length = length;
	rev->rev = this;

	if (is_tip()) name = only_child->name;

	delete only_child->next;
	delete only_child;

	return;
}


void TreeView::bottom_up(void (TreeView::* calc)()) 
{
	vector<TreeView*> c = childs();
	for(vector<TreeView*>::const_iterator i = c.begin(); i != c.end(); i++)
		(*i)->bottom_up(calc);

	(this->*calc)();
}

void TreeView::top_down(void (TreeView::* calc)())
{
	(this->*calc)();
	
	vector<TreeView*> p = parents();
	for(vector<TreeView*>::const_iterator i = p.begin(); i != p.end(); i++)
		(*i)->top_down(calc);	
}



vector<TreeView*> TreeView::childs()
{
	vector<TreeView*> out(0); // do I need that 0 to ensute it's empty?

	for(TreeView* t = rev->next; t != rev; t = t->next)
		out.push_back(t);

	return out;
}


vector<TreeView*> TreeView::parents()
{
	vector<TreeView*> out(0); // do I need that 0 to ensure it's empty?

	for(TreeView* t = next; t != this; t = t->next)
		out.push_back(t->rev);

	return out;
}

// This appends to the vector in input all the tips in the TreeView in input,
// and returns the number of tips appended.  Note that when I say tips in
// the TreeView, I'm excluding the tips in other TreeViews that may 
// include / intersect our TreeView (otherwise I would have said "tree").
//
int TreeView::get_tips(vector<TreeView*>& out)
{
	if (is_tip()) 
	{
		out.push_back(this);
		return 1;
	}

	int sum = 0;
	vector<TreeView*> c = childs();

	// if there are children, do the recursive calls:
	typedef vector<TreeView*>::const_iterator CI;
	for(CI i = c.begin(); i != c.end(); i++)
		sum += (*i)->get_tips(out);

	return sum;
}

// ########################### CATS FUNCTIONS #################################

void TreeView::calc_eplists() // need to add checks on memory usage !!!
{
	unsigned max_eplist_size_b4_clean    = 0;
	unsigned max_eplist_size_after_clean = 0;

	cout << "Subtree " << repr << flush;

	// Memory?
	eplist = vector<EPlist>(C+1); // A vector is created involving O(C)
	                              // initializations of empty EPlists and then
	                              // this vector is copied onto eplist [+O(C)].
	cout << " ->" << flush;

	if (is_tip()) 
		for(int c = 0; c <= C; c++)
			eplist[c].push_back(EPoint(length*(*pT)[c], (*pT)[c], c, 0, 0, this));
	else 
		for(int c = 0; c <= C; c++)
		{
			// eplist[c].reserve(...); ?
			for(int i = 0; i <= c; i++)             // O(cL^2)
				combine_into_eplist(c, rev->next->eplist[i], 
				             rev->next->next->eplist[c-i]);

			update_max(static_cast<unsigned>(eplist[c].size()), 
				max_eplist_size_b4_clean); 

			#ifndef DEBUG
			if (dNAP) clean_eplist(c);              // O(cL^2 * log cL)

			if (c != 0) cout << "\b\b\b\b\b\b";
			cout.width(6);
			cout << c << flush;
			update_max(static_cast<unsigned>(eplist[c].size()), 
				max_eplist_size_after_clean); 
			#endif
		}
	#ifdef DEBUG
	print_eplists("u");
	for(int c = 0; c <= C; c++)
	{
		cout << eplist[c].size() << "/";
		if (dNAP) clean_eplist(c);              // O(cL^2 * log cL)
		cout << eplist[c].size() << ", " << flush;
		update_max(static_cast<unsigned>(eplist[c].size()), 
			max_eplist_size_after_clean);
	}
	print_epinfo = true;
	#endif
	if (print_epinfo) print_eplists("t");
	
	cout << " done: " << max_eplist_size_b4_clean 
	     << "/" << max_eplist_size_after_clean << endl;
}


void TreeView::combine_into_eplist(int c, const EPlist& epl_L, const EPlist& epl_R)
{
	typedef EPlist::const_iterator CI;
	for (CI i = epl_L.begin(); i != epl_L.end(); i++)
		for (CI j = epl_R.begin(); j != epl_R.end(); j++)
		{
			// double p_new = i->p + j->p - (i->p * j->p);   // Sometimes this gives < 1 unduly...
			double p_new = 1.0 - (1.0 - i->p)*(1.0 - j->p);  // WHAT'S HAPPENING ???
			double e_new = i->e + j->e + p_new * length;

			// For the 0->1 NAP there will never be the need to store more than one
			// point in an Eplist:

			EPlist& epl_T = eplist[c];

			if (epl_T.size() >= 1 && 
				e_new <= epl_T.back().e &&
				p_new <= epl_T.back().p)
				continue;

			if (epl_T.size() >= 1 && 
				e_new >= epl_T.back().e &&
				p_new >= epl_T.back().p)
			{
				epl_T.back().e = e_new;
				epl_T.back().p = p_new;
				epl_T.back().left  = &(*i); // do I really have to do it like this?
				epl_T.back().right = &(*j);
				epl_T.back().c  =    c;
				epl_T.back().tv = this;
				continue;
			}
			// This may fail due to insufficient memory:
			epl_T.push_back(EPoint(e_new, p_new, c, &(*i), &(*j), this));
		}
}


void TreeView::clean_eplist(int c)  // Called O(nC) times.
{
	sort(eplist[c].begin(), eplist[c].end());

	EPlist vertices;

	// Find the index of the last point with maximum p:
	int top = 0;
	for(int i = 0; i < eplist[c].size(); i++)
		if (eplist[c][i].p >= eplist[c][top].p)
			top = i;

	// Go through eplist[c], and only push vertices onto 'vertices',
	// with a modified Graham's scan (CLRS pag.949):

	vertices.push_back(eplist[c][top]);
	if (top+1 < eplist[c].size())
		vertices.push_back(eplist[c][top+1]);

	for (int k, i = top+2; i < eplist[c].size(); i++)
	{
		while ((k = vertices.size()) >= 2 &&
			non_right_turn(vertices[k-2], vertices.back(), eplist[c][i]) >= 0)
				vertices.pop_back();
		vertices.push_back(eplist[c][i]);
	}
	eplist[c].swap(vertices);
}


EPoint* TreeView::get_best_epoint(int c)
{
	if (eplist[c].size() == 0) 
		throw string("launched get_best_epoint on an empty EPlist");

	EPoint* ep = &eplist[c][0];
	for (int i = 1; i < eplist[c].size(); i++)
		if (eplist[c][i].e > ep->e) ep = &eplist[c][i];
	return ep;
}	

// ########################## EPoint FUNCTIONS ################################

EPoint::EPoint(double e_new, double p_new, int c_new,
			   const EPoint* l, const EPoint* r, TreeView* tv_new)
{
	e = e_new;
	p = p_new;
	c = c_new;
	left  = l;
	right = r;
	tv = tv_new;
}

//
// This function adds to the input map a number of pairs (taxon,expenditure).
// Note that if this map already contains a number of pairs, possibly some of
// these will not be overwritten.
//
void EPoint::get_investments(map<string,int>& out) const
{
	if (left == 0 && right == 0)
		out[tv->get_name()] = c;
	else
	{
		left ->get_investments(out);
		right->get_investments(out);
	}
}


// ########################## OUTPUT FUNCTIONS ################################

void TreeView::print_eplists(char* prefix)
{
	// This assumes that the repr was calclated in the non-verbose mode:
	string filename = prefix + repr + ".eplist";

	ofstream out(filename.c_str());
	if (!out) throw string("cannot open " + filename + " for writing");
	out.precision(15);
	print_eplists(out);
}

void TreeView::print_eplists(ostream& out)
{
	for(int c = 0; c <= C; c++) 
		out << "e" << c << "\tp" << c << "\t";
	out << endl;

	int out_of_bounds = 0;
	for(int i = 0; out_of_bounds < C+1; i++)
	{
		out_of_bounds = 0;
		for(int c = 0; c <= C; c++)
			if (i >= eplist[c].size()) 
			{
				out.width(40);
				out << "-\t-\t";
				out_of_bounds++;
			}
			else out << eplist[c][i].e << '\t' << eplist[c][i].p << '\t';
			
		out << endl;
	}
}

//
// This function does a number of things. 
// It returns the width that the columns of the table will need to have.
//
int TreeView::set_header_of_output_table(ostream& out, vector<string>& taxon_names)
{
	vector<TreeView*> tips;
	get_tips(tips);

	int max_cost = 0;
	int max_name_length = 0;
	for(int i = 0; i < tips.size(); i++)
	{
		taxon_names.push_back(tips[i]->get_name());
		if (tips[i]->cost > max_cost) max_cost = tips[i]->cost;
		if (tips[i]->get_name().length() > max_name_length)
			max_name_length = tips[i]->get_name().length();
	}

	int col_width;
	if (max_cost == 0) col_width = 2;
	else col_width = log10(float(max_cost)) + 2;

	out << "Budget   Taxa";
	for (int k = 0; k < col_width * tips.size() - 4; k++)
		out << " ";
	if (dNAP) out << "    E(PD)" << endl;
	else out << "    PD" << endl;

	for(int i = 0; i < max_name_length; i++)
	{
		out << "         ";
		for(int j = 0; j < taxon_names.size(); j++)
		{
			out.width(col_width);
			if (i >= taxon_names[j].length()) out << " ";
			else out << taxon_names[j][i];
		}
		out << endl;
	}
	return col_width;
}

void TreeView::print_rPD_solutions_as_table(ostream& out)
{
	out.setf(ios::left);
	out << "Table of the solutions to the rooted problem" << endl;
    out << "--------------------------------------------" << endl;

	vector<string> taxon_names;
	int col_width = set_header_of_output_table(out, taxon_names);

	map<string, int> expenditures; // can define it here, as it will
	                               // always be completely overwritten

	EPoint* ep = 0;
	for(int budget = 0; budget <= C; budget++)
	{
		EPoint* ep_new = get_best_epoint(budget);
		if (ep == 0 || ep_new->e > ep->e)
		{
			out.width(9);
			out << budget;
			ep_new->get_investments(expenditures); // overwrites it completely

			for(int i = 0; i < taxon_names.size(); i++)
			{
				int x_i = expenditures[taxon_names[i]];

				out.width(col_width);
				if (x_i == 0) out << ".";
				else out << x_i;
			}
			out << "    " << (ep = ep_new)->e << endl;
		}
	}
	out << endl;
}

EPoint* try_and_improve(EPoint* ep, int budget, const vector<TreeView*>& tips,
						bool reduce_cost = false);

void TreeView::print_uPD_solutions_as_table(ostream& out)
{
	// This borrows a lot from other functions, and perhaps a better
	// organisation of the code is possible...

	if (dNAP) out << "Warning: printing unrooted-PD solution for dNAP!\n";

	out.setf(ios::left);
	out << "Table of the solutions to the unrooted problem" << endl;
    out << "----------------------------------------------" << endl;

	vector<string> taxon_names;
	int col_width = set_header_of_output_table(out, taxon_names);

	map<string, int> expenditures; // can define it here, as it will
	                               // always be completely overwritten

	vector<TreeView*> tips;
	get_tips(tips);
	rev->get_tips(tips);

	EPoint* old_ep = 0;

	for(int budget = 0; budget <= C; budget++)
	{
		EPoint* new_ep = old_ep;
		new_ep = try_and_improve(new_ep, budget, tips); // without reducing costs!!!

		int root_cost;

		if (new_ep == 0 || new_ep->e == 0)
		{
			new_ep = &(tips[0]->rev->eplist[0][0]); // Solution having 0 investments.
			root_cost = 0;
		}
		else root_cost = new_ep->tv->rev->cost;
		
		if (new_ep == old_ep) continue;

		out.width(9);
		out << budget;
		new_ep->get_investments(expenditures); // doesn't overwrite it completely

		expenditures[new_ep->tv->rev->get_name()] = root_cost;

		for(int i = 0; i < taxon_names.size(); i++)
		{
			int x_i = expenditures[taxon_names[i]];

			out.width(col_width);
			if (x_i == 0) out << ".";
			else out << x_i;
		}
		out << "    " << (old_ep = new_ep)->e << endl; 
	}
	out << endl;
}


//
// This must be called on the TreeView with respect to which the rooted PD
// is defined.
//
void TreeView::print_solution_for_rooted_PD(int budget, ostream& out)
{
	out << "----------------------------------------"
		   "---------------------------------------" << endl;
	out << "The rooted-PD optimal solution with budget " << budget << " is:" << endl;
	print_expenditures(out, reduced_cost_version( get_best_epoint(budget) ));
}

void TreeView::print_all_solutions_for_rooted_PD(ostream& out)
{
	EPoint* ep = 0;
	for(int budget = 0; budget <= C; budget++)
	{
		EPoint* ep_new = get_best_epoint(budget);
		if (ep == 0 || ep_new->e > ep->e)
		{
			out << "The rooted-PD optimal solution with budget " << budget << " is:" << endl;
			print_expenditures(out, ep = ep_new);
		}
	}
}

//
// This function behaves the same no matter which TreeView subtree of an
// unrooted tree it is called on.  I.e. it will output the solution relative 
// to the unrooted tree that 'this' is part of.  Time: O(nC) worst case.
//
void TreeView::print_solution_for_unrooted_PD(int budget, ostream& out) 
{
	if (dNAP) out << "Warning: printing unrooted-PD solution for dNAP!\n";

	vector<TreeView*> tips;
	get_tips(tips);
	rev->get_tips(tips);

	EPoint* ep = 0;
	ep = try_and_improve(ep, budget, tips, true);

	int root_cost;

	if (ep == 0 || ep->e == 0)
	{
		ep = &(tips[0]->rev->eplist[0][0]); // Solution having 0 investments.
		root_cost = 0;
	}
	else root_cost = ep->tv->rev->cost;

	out << "----------------------------------------"
		   "---------------------------------------" << endl;
	out << "The unrooted-PD optimal solution with budget " << budget << " is:" << endl;
	out.width(20);
	out << ep->tv->rev->name << "\t" << root_cost << endl;
	ep->tv->print_expenditures(out, ep, root_cost); 
}

void TreeView::print_all_solutions_for_unrooted_PD(ostream& out) 
{
	if (dNAP) out << "Warning: printing unrooted-PD solution for dNAP!\n";

	vector<TreeView*> tips;
	get_tips(tips);
	rev->get_tips(tips);

	EPoint* old_ep = 0;

	for(int budget = 0; budget <= C; budget++)
	{
		EPoint* ep = old_ep;
		ep = try_and_improve(ep, budget, tips);

		int root_cost;

		if (ep == 0 || ep->e == 0)
		{
			ep = &(tips[0]->rev->eplist[0][0]); // Solution having 0 investments.
			root_cost = 0;
		}
		else root_cost = ep->tv->rev->cost;
		
		if (ep == old_ep) continue;

		out << "The unrooted-PD optimal solution with budget " << budget << " is:" << endl;
		out.width(20);
		out << ep->tv->rev->name << "\t" << root_cost << endl;
		ep->tv->print_expenditures(out, old_ep = ep, root_cost); 
	}
}

EPoint* try_and_improve(EPoint* ep, int budget, const vector<TreeView*>& tips, 
						bool reduce_cost)
{
	for(int i = 0; i < tips.size(); i++)
	{
		if (tips[i]->cost > budget) continue;

		EPoint* tip_solution = tips[i]->get_rev().get_best_epoint(budget - tips[i]->cost);

		if (reduce_cost)
			tip_solution = tips[i]->get_rev().reduced_cost_version(tip_solution);

		if (ep == 0 || tip_solution->e > ep->e || 
			tip_solution->e == ep->e && tip_solution->c + tips[i]->cost < ep->c + ep->tv->get_rev().cost)
				ep = tip_solution;
	}
	return ep;
}


//
// This tries to reduce the cost of solution *ep, without decreasing its E[PD].
//
EPoint* TreeView::reduced_cost_version(EPoint* ep)
{
	EPoint *ep_new;
	while(ep->c >= 1 && (ep_new = get_best_epoint(ep->c - 1))->e == ep->e)
		ep = ep_new;
	return ep;
}

//
// This function is used by the print_solution_... functions.
// It prints out the expenditures associated with solution point *ep.
// The expenditure to assign to the root of the TreeView on which this function
// is called is given by 'root_cost'.  It's assumed that this expenditure has 
// already been printed.  Time: O(n).
//
void TreeView::print_expenditures(ostream& out, EPoint* ep, int root_cost)
{
	map<string,int> investments;
	ep->get_investments(investments);
	out.setf(ios::left,ios::adjustfield);

	typedef map<string,int>::const_iterator CI;
	for(CI i = investments.begin(); i != investments.end(); i++) if (i->first != "ROOT") 
	{
		out.width(20);
		out << i->first << "\t" << i->second << endl;
	}
	out << "------" << endl;
	out.width(20);
	out << "total:" << "\t" << ep->c + root_cost << endl;

	if (dNAP)
		out << "\nThis achieves E[PD] = " << ep->e;
	else
		out << "\nThis achieves PD = " << ep->e;

	out << "\n----------------------------------------"
		"---------------------------------------" << endl;
}



